import os,json
import numpy as np
import cv2
from util.sketlonize import Sketlonize
from util.corner_detecot import CornerDetector
import matplotlib.pyplot as plt
import math
def calculate_angle_diff(direction1, direction2):
    """
    Calculate the shortest difference between two angles.
    Takes into account the circular nature of angles (0-360 degrees).
    """
    diff = abs(direction1 - direction2)
    return min(diff, 360 - diff)
class VesselNode:
    """
    Represents a single vessel segment between key points.
    """

    def __init__(self, start, end, pixels,next=[]):
        self.start = start  # Start point (x, y)
        self.end = end  # End point (x, y)
        self.pixels = pixels  # List of all pixels in this segment
        self.direction = self.calculate_angle(start, end)  # Direction in degrees
        self.length = self.calculate_length(start, end)  # Length of the segment
        self.next =next  # List of next VesselNodes (for branches)
    def _boad_cast(self):
        print("vessel node")
        print(f"start:{self.start}")
        print(f"end:{self.end}")
        print(f"direction:{self.direction}")
        print(f"length:{self.length}")
        print(f"next:{self.next}")
    def __str__(self):
        return(f"start:{self.start},end:{self.end},direction:{self.direction},length:{self.length},next:{self.next}")
    @staticmethod
    def calculate_angle(start, end):
        """
        Calculate the angle (0-360 degrees) between the segment and the x-axis.
        """
        dx = end[0] - start[0]
        dy = end[1] - start[1]
        angle = math.degrees(math.atan2(dy, dx))
        return angle % 360

    @staticmethod
    def calculate_length(start, end):
        """
        Calculate the Euclidean distance between start and end.
        """
        return math.sqrt((end[0] - start[0])**2 + (end[1] - start[1])**2)


class Vessel:
    """
    Represents a complete vessel as a series of VesselNodes.
    """

    def __init__(self, start_point, nodes=[]):
        self.start_point = start_point  # Starting point of the vessel
        self.nodes = nodes  # List of VesselNodes

    def add_node(self, node):
        self.nodes.append(node)

    def compute_embedding(self):
        """
        Compute the embedding for the vessel as a list of direction and length vectors.
        """
        embedding = [(node.direction, node.length) for node in self.nodes]
        return embedding

class VesselVector:
    def __init__(self,direction_threshold=5,vessel_mini_length=3):
        self.direction_threshold=direction_threshold
        self.vessel_mini_length=vessel_mini_length
    def cal_point_distance(self,point1,point2):
        """
        Calculate the Euclidean distance between two points.
        """
        return math.sqrt((point1[0] - point2[0])**2 + (point1[1] - point2[1])**2)

    def MergeNode(self,node1,node2):
        """
        Merge two VesselNodes into a single VesselNode.
        """
        # check if the nodes are adjacent
        adj_dis=self.cal_point_distance(node1.end, node2.start)
        
        if adj_dis > 3:
            return False
        if not node1.next or len(node1.next)>1:
            # 不允许和分叉点前向合并
            return False
        start = node1.start
        end = node2.end
        pixels = node1.pixels + node2.pixels
        next=node2.next
        return VesselNode(start, end, pixels,next)
    
    def Simplify(self,vessel):
        simplified_nodes = []
        
        # 辅助变量
        last_direction=None
        last_node=None # 暂存和下一条血管合并的血管
        
        for idx,node in enumerate(vessel.nodes):
            if last_node:
                new_node=self.MergeNode(last_node,node)
                if new_node:
                    node=new_node
                else:
                    simplified_nodes.append(last_node)
                last_node=None
                # print("和之前存储的血管段的血管合并")
            # 是否需要和前段合并
            if last_direction and calculate_angle_diff(node.direction,last_direction)<self.direction_threshold and node.direction!=0:
                simplified_nodes[-1]=self.MergeNode(simplified_nodes[-1],node)
                # print("和已简化的合并")
                continue
            
            # 血管是否过短需要合并
            if node.length<self.vessel_mini_length:
                if idx==0 or len(simplified_nodes)==0:
                    # 第一段血管和后面的合并,
                    last_node=node
                    continue
                # 如果是分叉点，则只能前向合并或者最后一段
                if (node.next and len(node.next)>=2) or(node.next is None):
                    new_node=self.MergeNode(simplified_nodes[-1],node)
                    if new_node:
                        simplified_nodes[-1]=new_node
                    else:
                        simplified_nodes.append(node)
                    continue
                else:
                    # 是中间段落，则看和谁的变化更平滑
                    after_node= vessel.nodes[idx+1]
                    before_node= simplified_nodes[-1]
                    after_dir_change= calculate_angle_diff(after_node.direction,node.direction)
                    before_dir_change= calculate_angle_diff(before_node.direction,node.direction)
                    if after_dir_change>before_dir_change:
                        # 和后面的血管变化大，则前向合并
                        new_node= self.MergeNode(simplified_nodes[-1],node)
                        if new_node:
                            
                            simplified_nodes[-1]=self.MergeNode(simplified_nodes[-1],node)
                        else:
                            simplified_nodes.append(node)
                        continue
                    else:
                        # 后向合并
                        last_node=node
                        continue
            # 无需合并
            simplified_nodes.append(node)
            last_node=None
        print(f"Simplify vessel, Original nodes: {len(vessel.nodes)}, Simplified nodes: {len(simplified_nodes)}")
        new_vessel=Vessel(vessel.start_point,simplified_nodes)
        return new_vessel
    def important_point_detecot(self,vessel):
        """
        Detect important points in the vessel structure.
        """
        important_candidate=[]
        important_points=[]
        node_number=len(vessel.nodes)
        for idx in range(0,node_number-1):
            node=vessel.nodes[idx]
            # 分叉点直接加入
            if node.next and len(node.next)>=2:
                important_points.append(node)
                continue
            
            next_node=vessel.nodes[idx+1]
            direction_change=calculate_angle_diff(node.direction,next_node.direction)
            important_candidate.append((node,direction_change))
        # 选取变化最大的3个点
        important_candidate.sort(key=lambda x:x[1],reverse=True)
        important_points.extend([x[0] for x in important_candidate[:3]])
        return important_points
    def dir_rection_check(self, vessels, optic_disc_corr):
        """
        Check and correct the direction of vessel nodes relative to the optic disc.
    
        Parameters:
            vessels (list): List of Vessel objects.
            optic_disc_corr (tuple): (x, y) coordinates of the optic disc.
    
        Returns:
            List of corrected Vessel objects.
        """
        for vessel in vessels:
            for node in vessel.nodes:
                start = node.start
                end = node.end
    
                # Vector from the optic disc to the segment's start point
                vec_to_disc = (start[0] - optic_disc_corr[0], start[1] - optic_disc_corr[1])
                vec_direction = (end[0] - start[0], end[1] - start[1])  # Segment direction vector
    
                # Calculate the dot product to determine the angle
                dot_product = vec_to_disc[0] * vec_direction[0] + vec_to_disc[1] * vec_direction[1]
    
                if dot_product < 0:  # If the angle is acute, reverse the direction
                    node.start, node.end = node.end, node.start  # Swap start and end
                    node.direction = (node.direction + 180) % 360  # Update direction
        return vessels


    def visual_portant_points(self,sketlon,important_points,radius,save_path):
        """
        以radius为半径，重要点为圆心，画一个空心的黄色圆
        """
        img=sketlon.copy()
        # 如果是单通道图像，转换为三通道
        if len(img.shape)==2:
            img=cv2.cvtColor(img,cv2.COLOR_GRAY2BGR)
        for point in important_points:
            x,y=point.start
            cv2.circle(img,(x,y),radius,(0,255,255),thickness=1)
        cv2.imwrite(save_path,img)
        
                        
                    


class BloodVesselEmbedding:
    """
    Class for building vessel structures and embeddings from skeleton images.
    """

    def build_visited_matrix(self, skeleton, corner_points, bifurcation_points):
        """
        Create a visited matrix for the skeleton.
        """
        
        visited = np.zeros_like(skeleton, dtype=np.uint8)
        visited[skeleton > 0] = 1  # Mark all skeleton pixels as unvisited
        for x, y in corner_points:  # Points are in (x, y) format
            visited[y, x] = 2  # Mark corners
        for x, y in bifurcation_points:  # Points are in (x, y) format
            visited[y, x] = 3  # Mark bifurcations
        return visited

    def traverse_segment(self, current, visited):
        """
        Traverse a vessel segment from a starting point until a key point (corner or bifurcation).
        """
        segment_pixels = [current]
        x, y = current

        while True:
            visited[y, x] = 4  # Mark as visited
            neighbors = self.get_neighbors(x, y, visited)
            if not neighbors:
                break
            # if multi neighbors
                
            selected_neighbor = neighbors[0]
            # If  reaches a key point, stop traversal
            if visited[selected_neighbor[1], selected_neighbor[0]] ==2:
                # 拐点
                return segment_pixels, [selected_neighbor]
            elif visited[selected_neighbor[1], selected_neighbor[0]] ==3:
                # 分叉点
                visited[selected_neighbor[1], selected_neighbor[0]] = 4
                segment_pixels.append(selected_neighbor)
                next_branch=self.get_neighbors(selected_neighbor[0], selected_neighbor[1], visited)
                
                return segment_pixels, next_branch
            # Continue along the vessel
            x, y = selected_neighbor
            segment_pixels.append((x, y))

        return segment_pixels, None
    def check_connected(self,mask,point1,point2):
        '''
        检查连通性质，查看两点之间是否存在可达的通路
        '''
        pass
    def traverse_vessel(self, start_point, visited,idx=0):
        """
        Traverse a full vessel starting from a given point.
        """
        stack = [start_point]
        vessel=Vessel(start_point,[])
        while stack:
            
            current = stack.pop()
            x, y = current
            if visited[y, x] == 4:
                continue
            segment_pixels, next_points = self.traverse_segment(current, visited)
            # Create a VesselNode and add to the vessel
            node = VesselNode(segment_pixels[0],segment_pixels[-1] , segment_pixels,next_points)
            # node._boad_cast()
            vessel.add_node(node)
            if next_points is None:
                continue
            # Add next point to stack for further traversal
            stack.extend(next_points)
        return vessel

    def get_neighbors(self, x, y, visited):
        """
        Get the valid neighbors of a pixel that are part of the skeleton.
        """
        neighbors = [
           (x - 1, y), (x, y - 1), (x, y + 1),(x + 1, y), # 先考虑上下左右
            (x + 1,y - 1), (x + 1, y + 1), (x - 1, y - 1), (x - 1, y + 1),#  再考虑对角线
        ]
        return [(nx, ny) for nx, ny in neighbors if 0 <= ny < visited.shape[0] and 0 <= nx < visited.shape[1] and visited[ny, nx] in [1, 2, 3]]
    
    def build_vessel_embedding(self, skeleton, vascular_cnt, vascular_cnt_img, corner_points, bifurcation_points, optic_disc_coor):
        """
        Build vessel embeddings for all vessels in the skeleton.

        Parameters:
            skeleton (np.ndarray): Binary skeleton image.
            vascular_cnt (int): Number of connected regions in the vascular_cnt_img.
            vascular_cnt_img (np.ndarray): H*W array, where each region has a unique label.
            corner_points (list of tuples): List of corner points (x, y).
            bifurcation_points (list of tuples): List of bifurcation points (x, y).
            optic_disc_coor (tuple): Coordinate (x, y) of the optic disc.

        Returns:
            vessels (list): List of Vessel objects representing all vessels.
        """
        visited = self.build_visited_matrix(skeleton, corner_points, bifurcation_points)

        # Initialize a list to store the starting point for each vascular region
        start_point_list = []

        for label in range(1, vascular_cnt+1):
            # Get all pixels belonging to the current vascular region
            region_pixels = np.argwhere(vascular_cnt_img == label)
            if len(region_pixels) == 0:
                raise ValueError(f"No pixels found for label {label}")
            # Find the pixel closest to the optic disc coordinate
            closest_point = min(region_pixels, key=lambda p: np.linalg.norm((p[1] - optic_disc_coor[0], p[0] - optic_disc_coor[1])))
            start_point_list.append((closest_point[1], closest_point[0]))  # Append in (x, y) format
        self.visual_visited(visited,[[307, 262]],f'./experiments/new_feature/debug.png')
        # Traverse each vascular region starting from its closest point
        vessels = []
        for i,label in enumerate(range(1, vascular_cnt + 1)):
            vessel=self.traverse_vessel(start_point_list[label - 1], visited,i)
            vessels.append(vessel)
        
        return vessels
    
    def visual_visited(self, visited, points, save_path):
        # 分别用 蓝 绿 红 黄 表示 未访问，拐点，分叉点，已访问
        visited_img = np.zeros((visited.shape[0], visited.shape[1], 3), dtype=np.uint8)
        visited_img[visited == 1] = [255, 0, 0]      # 未访问 - 红色
        visited_img[visited == 2] = [0, 255, 0]      # 拐点 - 绿色
        visited_img[visited == 3] = [0, 0, 255]      # 分叉点 - 蓝色
        visited_img[visited == 4] = [0, 255, 255]    # 已访问 - 黄色

        # 在 points 的位置画一个半径为3的实心黄色圆圈
        for (x, y) in points:
            cv2.circle(visited_img, (x, y), radius=3, color=(0, 255, 255), thickness=-1)  # 黄色，实心圆

        # 保存图片
        cv2.imwrite(save_path, visited_img)

    def visualize_vessels(self, skeleton, vessels, corner_points, bifurcation_points, save_path=None):
        """
        Visualize vessels with alternating colors and highlight key points.
        """
        # Create an RGB image for visualization
        height, width = skeleton.shape
        visualization = np.zeros((height, width, 3), dtype=np.uint8)

        # Define a fixed color palette
        color_palette = [
            (255, 0, 0),    # Red
            (255, 255, 0),  # Yellow
            (0, 0, 255),    # Blue
            (0, 255, 0)     # Green
        ]
        # Visualize vessels with alternating colors
        for idx, vessel in enumerate(vessels):
            for j,node in enumerate(vessel.nodes):
                color = color_palette[j% len(color_palette)]  # Cycle through the color palette
                for pixel in node.pixels:
                    x, y = pixel
                    visualization[y, x] = color

        # Optional: Highlight corner points and bifurcation points
        # for x, y in corner_points:
        #     cv2.circle(visualization, (x, y), radius=3, color=(255, 255, 255), thickness=-1)  # White for corner points

        # for x, y in bifurcation_points:
        #     cv2.circle(visualization, (x, y), radius=3, color=(0, 255, 255), thickness=-1)  # Cyan for bifurcation points

        # Display or save the result
        plt.figure(figsize=(10, 10))
        plt.imshow(visualization)
        plt.title("Blood Vessel Visualization")
        plt.axis('off')

        if save_path:
            plt.savefig(save_path)
    def visual_vessel_direction(self, skeleton, vessels, save_path):
        """
        Visualize vessel directions by drawing arrows from node.start to node.end.
        """
        # Create an RGB image for visualization
        img = skeleton.copy()
        if len(img.shape) == 2:  # If single-channel, convert to 3-channel
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

        # Arrow drawing parameters
        arrow_color = (0, 255, 0)  # Green
        arrow_thickness = 1

        for vessel in vessels:
            for node in vessel.nodes:
                start = (int(node.start[0]), int(node.start[1]))
                end = (int(node.end[0]), int(node.end[1]))
                # Draw an arrowed line from start to end
                cv2.arrowedLine(
                    img, start, end, arrow_color, thickness=arrow_thickness, tipLength=0.3
                )

        # Save the resulting image
        cv2.imwrite(save_path, img)
        
def build_vessels(vascular_path,optic_disc_coor,com_press_r=3/8,detect_important_point=False):
    skeletoner=Sketlonize(0.35)
    x,y=optic_disc_coor
    x=int(x*com_press_r)
    y=int(y*com_press_r)
    skeleton,vascular_cnt_img,vascular_cnt=skeletoner.sketlonize(vascular_path,optic_disc_coor=(x,y), min_size=100)
    # 拐点检测
    corner_detector = CornerDetector()
    points = corner_detector.point_detector(skeleton)
    # corner_detector.visual_point(skeleton, points, save_path=os.path.join(visual_dir,'harris.jpg'))

    # 分叉点检测
    bifurcation_points = corner_detector.detect_bifurcation_points(skeleton,points)
    
    # 创建 BloodVesselEmbeddingWithVisualization 实例
    vessel_embedding = BloodVesselEmbedding()
    
    # 构建血管结构和嵌入
    vessels = vessel_embedding.build_vessel_embedding(skeleton, vascular_cnt, vascular_cnt_img, points, bifurcation_points, (x, y))
    
    vessel_vector=VesselVector()
    simplied_vessels=[]
    important_points=[]
    for vessel in vessels:
        simply_vessel=vessel_vector.Simplify(vessel)
        simplied_vessels.append(simply_vessel)
        important_points.extend(vessel_vector.important_point_detecot(simply_vessel))
        
    # 纠正血管方向
    vessels = vessel_vector.dir_rection_check(simplied_vessels, (x, y))
    
    if detect_important_point:
        return vessels,important_points
    
    return vessels
if __name__ == '__main__':

    # 获取配置
    data_path='../Dataset/infantImages'
    # 读取数据字典
    with open(os.path.join(data_path, 'annotations.json'), 'r') as f:
        data_dict = json.load(f)
    visual_dir='./experiments/new_feature'
    if not os.path.exists(visual_dir):
        os.makedirs(visual_dir)
    data_set_path='./experiments/dataset'
    vascular_dir=os.path.join(data_set_path,'vascular')
    image_dir= os.path.join(data_set_path,'images')
    image_name= '1243.jpg'
    vascular_path=os.path.join(vascular_dir,image_name)
    skeletoner=Sketlonize(0.35)
    x,y=data_dict[image_name]['optic_disc_pred']['position']
    x=int(x*3/8)
    y=int(y*3/8)
    
    skeleton,vascular_cnt_img,vascular_cnt=skeletoner.sketlonize(vascular_path,optic_disc_coor=(x,y), min_size=100)
    # 拐点检测
    corner_detector = CornerDetector()
    points = corner_detector.point_detector(skeleton)
    # corner_detector.visual_point(skeleton, points, save_path=os.path.join(visual_dir,'harris.jpg'))

    # 分叉点检测
    bifurcation_points = corner_detector.detect_bifurcation_points(skeleton,points)
    
    # 创建 BloodVesselEmbeddingWithVisualization 实例
    vessel_embedding = BloodVesselEmbedding()
    
    # 构建血管结构和嵌入
    vessels = vessel_embedding.build_vessel_embedding(skeleton, vascular_cnt, vascular_cnt_img, points, bifurcation_points, (x, y))
    
    # 可视化结果
    vessel_embedding.visualize_vessels(skeleton, vessels, points, bifurcation_points, save_path=os.path.join(visual_dir, 'vessels.jpg'))
    
    # 重要点检测
    vessel_vector=VesselVector()
    simplied_vessels=[]
    important_points=[]
    for vessel in vessels:
        simply_vessel=vessel_vector.Simplify(vessel)
        simplied_vessels.append(simply_vessel)
        important_points.extend(vessel_vector.important_point_detecot(simply_vessel))
    
    # 可视化重要点
    vessel_vector.visual_portant_points(skeleton,important_points,5,os.path.join(visual_dir,'important_points.jpg'))
    
    # 可视化简化后的血管
    vessel_embedding.visualize_vessels(skeleton, simplied_vessels, points, bifurcation_points, save_path=os.path.join(visual_dir, 'simplied_vessels.jpg'))
    
    # 纠正血管方向
    vessels = vessel_vector.dir_rection_check(simplied_vessels, (x, y))
    # 可视化血管方向
    vessel_embedding.visual_vessel_direction(skeleton, vessels, save_path=os.path.join(visual_dir, 'vessel_direction.jpg'))